#include "socket_helper.h"
#include <errno.h>
#include <pthread.h>
#include <netdb.h>
#include <netinet/in.h>
#include <unistd.h>
#include "proxy.h"
#include <sys/socket.h>
#include <stdlib.h>
#include "debug.h"
#include "util.h"

static int touch_header_end(char* line_buffer) {
	if (strncmp(line_buffer, "\r\n", 2) == 0)
		return 1;
	else
		return 0;
}
ssize_t readLine(int fd, void* buffer, size_t n) {
	ssize_t numRead;
	size_t  totRead;
	char*   buf;
	char    ch;

	if (n <= 0 || buffer == NULL) {
		errno = EINVAL;
		return -1;
	}

	buf = buffer;

	totRead = 0;
	for (;;) {
		numRead = recv(fd, &ch, 1, 0);

		if (numRead == -1) {
			if (errno == EINTR)
				continue;
			else
				return -1; /* 未知错误 */
		}
		else if (numRead == 0) {  /* EOF */
			if (totRead == 0) /* No bytes read; return 0 */
				return 0;
			else /* Some bytes read; add '\0' */
				break;
		}
		else {

			if (totRead < n - 1) { /* Discard > (n - 1) bytes */
				totRead++;
				*buf++ = ch;
			}

			if (ch == '\n')
				break;
		}
	}

	*buf = '\0';
	return totRead;
}
static char* read_header(int sock) {
	char*	header_buff    = NULL;
	unsigned int MAX_HEADER_LEN = 9000;
	char	 line_buffer[ 2048 ];
	char*	base_ptr;

	header_buff = ( char* )calloc(MAX_HEADER_LEN, sizeof(char));
	if (NULL == header_buff) {
		LOG("can't calloc to head_buf in read_header()\n");
		return NULL;
	}
	base_ptr = header_buff;

	while (1) {
		memset(line_buffer, 0, 2048);

		int total_read = readLine(sock, line_buffer, 2048);
		if (total_read <= 0) {
			return NULL;
		}

		if (base_ptr + total_read - header_buff > MAX_HEADER_LEN) {

			MAX_HEADER_LEN += 2048;
			header_buff = realloc(header_buff, MAX_HEADER_LEN);
			if (NULL == header_buff) {
				LOG("realloc to header_buff failed \n");
				return NULL;
			}
		}

		strncpy(base_ptr, line_buffer, total_read);
		base_ptr += total_read;

		if (touch_header_end(line_buffer))
			break;
		//读到了空行，http头结束
		if (strcmp(line_buffer, "\r\n") == 0
		    || strcmp(line_buffer, "\n") == 0) {
			break;
		}
	}

	return header_buff;
}

int extract_host_and_port(const char* header, hp_t* hp) {

	memset(hp->host, 0, sizeof(hp->host));

	char* _p = strstr(
		header,
		"CONNECT"); /* 在 CONNECT 方法中解析 隧道主机名称及端口号 */
	if (_p) {
		char* _p1 = strchr(_p, ' ');

		char* _p2 = strchr(_p1 + 1, ':');
		char* _p3 = strchr(_p1 + 1, ' ');

		if (_p2) {
			char s_port[ 10 ];
			bzero(s_port, 10);

			strncpy(hp->host, _p1 + 1, ( int )(_p2 - _p1) - 1);
			strncpy(s_port, _p2 + 1, ( int )(_p3 - _p2) - 1);
			hp->port = atoi(s_port);
		}
		else {
			strncpy(hp->host, _p1 + 1, ( int )(_p3 - _p1) - 1);
			hp->port = 80;
		}

		return 0;
	}

	char* p = strstr(header, "Host:");
	if (!p) {
		return -1;
	}
	char* p1 = strchr(p, '\n');
	if (!p1) {
		return -2;
	}

	char* p2 = strchr(p + 5, ':'); /* 5是指'Host:'的长度 */

	if (p2 && p2 < p1) {

		int  p_len = ( int )(p1 - p2 - 1);
		char s_port[ p_len ];
		strncpy(s_port, p2 + 1, p_len);
		s_port[ p_len ] = '\0';
		hp->port	= atoi(s_port);

		int h_len = ( int )(p2 - p - 5 - 1);
		strncpy(hp->host, p + 5 + 1, h_len);  // Host:
		// assert h_len < 128;
		hp->host[ h_len ] = '\0';
	}
	else {
		int h_len = ( int )(p1 - p - 5 - 1 - 1);
		strncpy(hp->host, p + 5 + 1, h_len);
		// assert h_len < 128;
		hp->host[ h_len ] = '\0';
		hp->port	  = 80;
	}
	return 0;
}
static int is_http_tunnel(const char* header) {
	if (strstr(header, "CONNECT") != 0)
		return 1;
	else
		return 0;
}
int report_tunnel_established(int client_sock) {
	char* resp = "HTTP/1.1 200 Connection Established\r\n\r\n";
	int   len  = strlen(resp);
	if (send(client_sock, resp, len, 0) < 0) {
		perror("Send http tunnel response  failed\n");
		return -1;
	}
	return 0;
}

void* thread_handler_push_data(void* arg) {
	sock_pair_t*       sock_pair = ( sock_pair_t* )arg;
	const unsigned int BUF_SIZE  = 8192;
	char*		   buffer    = ( char* )calloc(BUF_SIZE, sizeof(char));
	int		   n;

	while ((n = recv(sock_pair->sock_src, buffer, BUF_SIZE, 0))
	       > 0) {
		send(sock_pair->sock_dest, buffer, n, 0);
	}
	destroy_mem(buffer);
	destroy_mem(sock_pair);

	LOG("push done\n");
	return NULL;
}

static void forward_data(int socka, int sockb) {
	pthread_t t1, t2;

	sock_pair_t *sock_pair_atob, *sock_pair_btoa;
	sock_pair_atob		  = ( sock_pair_t* )malloc(sizeof(sock_pair_t));
	sock_pair_btoa		  = ( sock_pair_t* )malloc(sizeof(sock_pair_t));
	sock_pair_atob->sock_src  = socka;
	sock_pair_atob->sock_dest = sockb;
	sock_pair_btoa->sock_dest = socka;
	sock_pair_btoa->sock_src  = sockb;

	pthread_create(&t1, NULL, ( void* )thread_handler_push_data,
		       sock_pair_atob);
	pthread_create(&t2, NULL, ( void* )thread_handler_push_data,
		       sock_pair_btoa);

	pthread_join(t1, NULL);
	pthread_join(t2, NULL);

	close(socka);
	close(sockb);
}
static void https_proxy(int client_sock, int remote_sock) {

	report_tunnel_established(client_sock);
	forward_data(client_sock, remote_sock);
}
static void rewrite_path_in_header(char *header_buffer) {
	char* p   = strstr(header_buffer, "http://");
	char* p0  = strchr(p, '\0');
	char* p5  = strstr(header_buffer,
			   "HTTP/"); /* "HTTP/" 是协议标识 如 "HTTP/1.1" */
	int   len = strlen(header_buffer);
	if (p) {
		char* p1 = strchr(p + 7, '/');
		if (p1 && (p5 > p1)) {
			//转换url到 path
			memcpy(p, p1, ( int )(p0 - p1));
			int l		   = len - (p1 - p);
			header_buffer[ l ] = '\0';
		}
		else {
			char* p2 = strchr(
				p, ' ');  // GET http://3g.sina.com.cn HTTP/1.1

			// printf("%s\n",p2);
			memcpy(p + 1, p2, ( int )(p0 - p2));
			*p		   = '/';  // url 没有路径使用根
			int l		   = len - (p2 - p) + 1;
			header_buffer[ l ] = '\0';
		}
	}
}
static void http_proxy(int client_sock, int remote_sock, char* header) {
	rewrite_path_in_header(header);	
	send(remote_sock, header, strlen(header), 0);
	LOG("http header : \n%s\n",header);
	forward_data(client_sock, remote_sock);
}
static void handle_client(int client_sock) {

	char* header;
	hp_t  remote_hp;

	LOG("============handle a new client==============\n");

	if (NULL == (header = read_header(client_sock))) {
		LOG("can't get header\n");
		exit(1);
	}

	if (extract_host_and_port(header, &remote_hp) < 0) {
		LOG("can't get host and port\n");
		exit(1);
	}

	int remote_sock = connect_with_remote(remote_hp.host, remote_hp.port);
	if (remote_sock < 0) {
		LOG("can't connect with remote %s:%d\n", remote_hp.host,
		    remote_hp.port);
		exit(1);
	}
	LOG("connet with remote %s:%d\n", remote_hp.host, remote_hp.port);

	if (is_http_tunnel(header)) 
		https_proxy(client_sock, remote_sock);
	else 
		http_proxy(client_sock, remote_sock, header);

	destroy_mem(header);
	LOG("==============================================\n");
}

void run_proxy(int local_sock) {

	int client_sock = 0;

	while (1) {
		client_sock = accept(local_sock, NULL, NULL);

		if (fork() == 0) {
			close(local_sock);
			handle_client(client_sock);
			exit(0);
		}

		close(client_sock);
	}
}